{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "metadata": {}
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoTokenizer\n",
    "from datasets import load_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "metadata": {}
   },
   "outputs": [],
   "source": [
    "model_name = \"meta-llama/Meta-Llama-3-8B-Instruct\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['messages'],\n",
       "    num_rows: 101476\n",
       "})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = load_dataset(\"sarvamai/samvaad-hi-v1\", split=\"train\")\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>messages</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[{'content': 'Can you tell me about the 1988 R...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[{'content': 'Can you tell me more about the m...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>[{'content': 'Mujhe Lavasa ke baare mein jaank...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>[{'content': 'वैश्वीकरण भारत में स्वदेशी संस्क...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>[{'content': 'Could you please tell me in Hind...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>[{'content': 'मुझे एलामकुलम के बारे में ज्यादा...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>[{'content': 'मेरा गुजारिश है कि आप हिंदी में ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>[{'content': 'मुझे जानकारी चाहिए करम बठ जी के ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>[{'content': 'What are the Schouten Islands?',...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>[{'content': 'Tell me more about the Deodhar T...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            messages\n",
       "0  [{'content': 'Can you tell me about the 1988 R...\n",
       "1  [{'content': 'Can you tell me more about the m...\n",
       "2  [{'content': 'Mujhe Lavasa ke baare mein jaank...\n",
       "3  [{'content': 'वैश्वीकरण भारत में स्वदेशी संस्क...\n",
       "4  [{'content': 'Could you please tell me in Hind...\n",
       "5  [{'content': 'मुझे एलामकुलम के बारे में ज्यादा...\n",
       "6  [{'content': 'मेरा गुजारिश है कि आप हिंदी में ...\n",
       "7  [{'content': 'मुझे जानकारी चाहिए करम बठ जी के ...\n",
       "8  [{'content': 'What are the Schouten Islands?',...\n",
       "9  [{'content': 'Tell me more about the Deodhar T..."
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = dataset.to_pandas()\n",
    "df.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "metadata": {}
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# data_point = [{'content': 'Can you tell me about the 1988 Rajya Sabha elections in English please?', 'role': 'user'}, {'content': \"The 1988 Rajya Sabha elections were a significant event in Indian politics. The Rajya Sabha, which is the upper house of India's Parliament, sees periodic elections with members serving terms of six years. In 1988, elections took place in several states across India, electing new members to represent their respective constituencies.\\n\\nThe elected members were expected to serve from 1988 to 1994. However, there was a possibility for their term to end sooner due to circumstances such as resignation or death.\\n\\nAlso, apart from the main elections, specific special elections known as 'bye elections' also occurred in 1988. These are typically held to fill vacancies that arise during the tenure of the elected members.\", 'role': 'assistant'}, {'content': 'Was it a norm for members to serve for a term of six years back then?', 'role': 'user'}, {'content': \"Yes, the norm for the Rajya Sabha members was and continues to be to serve for a term of six years. This tenure is a standard part of India's democratic process and applies nationwide, regardless of the member's state or party affiliation. After the term ends, the members could choose to stand for re-election. Please note that although the term is for six years, a member's tenure may end before this period if they resign, pass away, or are otherwise unable to fulfill their obligations.\", 'role': 'assistant'}, {'content': 'What is the purpose of the bye-elections?', 'role': 'user'}, {'content': 'By-elections, also known as special elections, are held to fill any political office that becomes vacant between the regular scheduled elections. The reasons for these vacancies can vary and include situations such as the death of an existing member, the resignation of a member, or the invalidation of an election.\\n\\nIn the context of the 1988 Rajya Sabha elections, bye-elections took place to fill the seats of any member who could not continue their term until 1994. The contestants for these bye-elections vary, and could include defeated candidates from the main election or new candidates from the respective political parties. This ensures that the democratic representation in the Rajya Sabha continues uninterrupted.', 'role': 'assistant'}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing dataset:   0%|          | 0/101476 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing dataset: 100%|██████████| 101476/101476 [00:10<00:00, 9494.11it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['messages', 'llama3_prompt_format', 'text', 'prompt_template'],\n",
       "    num_rows: 101476\n",
       "})"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "from tqdm import tqdm\n",
    "\n",
    "def json_to_prompt(json_data):\n",
    "    messages = [{\"role\": \"system\", \"content\": \"You are Gaja, an AI assistant created by Cognitivelab and trained on top of Llama 3 Large language model(LLM), proficient in English and Hindi. You can respond in both languages based on the users request.\"}]\n",
    "        \n",
    "    # messages = [\n",
    "    #     {\"role\": \"system\", \"content\": \"You are a pirate chatbot who always responds in pirate speak!\"},\n",
    "    #     {\"role\": \"user\", \"content\": \"Who are you?\"},\n",
    "    # ]\n",
    "    for entry in json_data:\n",
    "        # print(entry)\n",
    "        content = entry[\"content\"]\n",
    "        role = entry[\"role\"]\n",
    "        if role == \"user\":\n",
    "            formatted_json = {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": content\n",
    "            }\n",
    "        #     formatted_string += f\"<start_of_turn>user {content} <end_of_turn>\\n\"\n",
    "        elif role == \"assistant\":\n",
    "            formatted_json = {\n",
    "                \"role\": \"assistant\",\n",
    "                \"content\": content\n",
    "            }\n",
    "        #     formatted_string += f\"<start_of_turn>model {content} <end_of_turn>\\n\"\n",
    "        messages.append(formatted_json)\n",
    "    \n",
    "    # print(messages)\n",
    "    tokenized_message = tokenizer.apply_chat_template(\n",
    "        messages, \n",
    "        tokenize=False, \n",
    "        add_generation_prompt=False\n",
    "    )\n",
    "    # print(tokenized_message)\n",
    "    \n",
    "    chat_template = messages\n",
    "\n",
    "    return tokenized_message, chat_template\n",
    "\n",
    "\n",
    "def generate_prompt(data_point):\n",
    "    \"\"\"Gen. input text based on a prompt, task instruction, (context info.), and answer\n",
    "\n",
    "    :param data_point: dict: Data point\n",
    "    :return: dict: tokenzed prompt\n",
    "    \"\"\"\n",
    "    \n",
    "    # print(data_point)\n",
    "    # json_data_parsed =  json.loads(data_point)\n",
    "    # text,json = json_to_prompt(data_point)\n",
    "    text,json = json_to_prompt(data_point[\"messages\"])\n",
    "    # text= json_to_prompt(data_point[\"messages\"])\n",
    "\n",
    "    return text,json\n",
    "\n",
    "\n",
    "# print(generate_prompt(data_point))\n",
    "\n",
    "# # add the \"prompt\" column in the dataset\n",
    "# text_column,message_column = [generate_prompt(data_point) for data_point in dataset]\n",
    "# dataset = dataset.add_column(\"llama3_prompt_formate\", text_column)\n",
    "# dataset = dataset.add_column(\"text\", text_column)\n",
    "# dataset = dataset.add_column(\"prompt_template\", message_column)\n",
    "\n",
    "# Create lists to store results\n",
    "text_column = []\n",
    "message_column = []\n",
    "\n",
    "# Process each data point in the dataset\n",
    "for data_point in tqdm(dataset, desc=\"Processing dataset\"):\n",
    "    # Generate chat prompt\n",
    "    text,chat_template = generate_prompt(data_point)\n",
    "    \n",
    "    # Append results to respective lists\n",
    "    text_column.append(text)\n",
    "    message_column.append(chat_template)\n",
    "\n",
    "# Add columns to the dataset\n",
    "dataset = dataset.add_column(\"llama3_prompt_format\", text_column)\n",
    "dataset = dataset.add_column(\"text\", text_column)\n",
    "dataset = dataset.add_column(\"prompt_template\", message_column)\n",
    "\n",
    "# Display updated dataset\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1da2d5a97e7144518fe46bb443cd34dc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading the dataset shards:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "175ea0cd923e4c399ddd2999dd9672ad",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/21 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "62b7c6d2a1f241b29d6e08298945c866",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/21 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9326e403bea04ba2b76b88d02827a2d5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/21 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "09d12afe11cc471daef023a6bca99fdf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/21 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f0910e6226c743419e564ca33d81d468",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/21 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/datasets/CognitiveLab/Samvaad_Hindi_Hinglish_Llama3_Prompt_formate/commit/51efdff6dcf929cc88143c3cd675286fbb75dc29', commit_message='Upload dataset', commit_description='', oid='51efdff6dcf929cc88143c3cd675286fbb75dc29', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.push_to_hub(\"CognitiveLab/Samvaad_Hindi_Hinglish_Llama3_Prompt_formate\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "metadata": {}
   },
   "outputs": [],
   "source": [
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": \"You are a pirate chatbot who always responds in pirate speak!\"},\n",
    "    {\"role\": \"user\", \"content\": \"Who are you?\"},\n",
    "    {\"role\": \"assistant\", \"content\": \"I am llama3 the best LLM out there\"}\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "metadata": {}
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<|begin_of_text|><|start_header_id|>system<|end_header_id|>\\n\\nYou are a pirate chatbot who always responds in pirate speak!<|eot_id|><|start_header_id|>user<|end_header_id|>\\n\\nWho are you?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\nI am llama3 the best LLM out there<|eot_id|>'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenized = tokenizer.apply_chat_template(\n",
    "        messages, \n",
    "        tokenize=False, \n",
    "        add_generation_prompt=False\n",
    ")\n",
    "\n",
    "tokenized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "LLM-venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
